from langchain_core.messages import SystemMessage, trim_messages,HumanMessage,AIMessage

# 只统计一下字母数量当作token数量
def token_counter(messages):
    total = 0
    for msg in messages:
        total += len(msg.content)
    return total


trimmer = trim_messages(
    max_tokens=28, # 最大token数量
    token_counter=token_counter, # 计算token数量的方法
    strategy="last", # 选择裁剪策略 "first"代表新加入的消息裁剪掉, "last"代表前面的消息裁剪掉
    include_system=False, # 裁剪完的消息是否包含系统消息
    allow_partial=True, # 是否允许一段message内携带的内容被裁剪掉
    start_on="human", # 从什么角色的message保留,如果设置human则，起始位置是human
)


messages = [
    SystemMessage(content="你是一个助手"),
    HumanMessage(content="我叫鲍勃"),
    AIMessage(content="你好"),
    HumanMessage(content="我喜欢冰淇淋"),
    AIMessage(content="很好"),
    HumanMessage(content="2+2等于几"),
    AIMessage(content="4"),
    HumanMessage(content="谢谢"),
    AIMessage(content="不客气"),
    HumanMessage(content="开心吗？"),
    AIMessage(content="是的！"),
]

r = trimmer.invoke(messages)
for message in r:
    print(message.content)